{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 174,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dotenv import load_dotenv\n",
    "\n",
    "# Load environment variables (such as API keys) from config.env\n",
    "load_dotenv('./config.env')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_ollama import ChatOllama\n",
    "\n",
    "model = ChatOllama(model=\"qwen2.5:7b\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.tools import tool\n",
    "from typing import Annotated\n",
    "@tool\n",
    "def calculate_fire_number(annual_expenses: Annotated[float, \"The person's expected annual expenses in retirement.\"]) -> float:\n",
    "    \"\"\"\n",
    "    This function calculates the FIRE number, given a person's expected annual expenses in retirement.\n",
    "    This is done by multiplying the annual expenses by 25.\n",
    "    \"\"\"\n",
    "    return 25 * annual_expenses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_with_tools = model.bind_tools([calculate_fire_number])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'langchain_core.messages.ai.AIMessage'>\n",
      "content='' additional_kwargs={} response_metadata={'model': 'qwen2.5:7b', 'created_at': '2025-02-16T08:26:28.33274Z', 'done': True, 'done_reason': 'stop', 'total_duration': 12856702666, 'load_duration': 32561708, 'prompt_eval_count': 222, 'prompt_eval_duration': 5977000000, 'eval_count': 78, 'eval_duration': 6822000000, 'message': Message(role='assistant', content='', images=None, tool_calls=None)} id='run-9c809417-a982-48f9-9adb-1a7da8b1c511-0' tool_calls=[{'name': 'calculate_fire_number', 'args': {'annual_expenses': 10000}, 'id': '1010a84e-5550-4a80-89cd-ff36b15c60d7', 'type': 'tool_call'}] usage_metadata={'input_tokens': 222, 'output_tokens': 78, 'total_tokens': 300}\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.messages import HumanMessage, SystemMessage\n",
    "\n",
    "messages = [SystemMessage(\"Do not hallucinate or assume any information. For example, if annual expenses are not provided, ask the user for it. After getting annual expenses from the user, call the tool 'calculate_fire_number' to get the FIRE number for the user.\"),\n",
    "            HumanMessage(\"What is my FIRE number? 10K.\")]\n",
    "\n",
    "\n",
    "response = model_with_tools.invoke(messages)\n",
    "\n",
    "print(type(response))\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[SystemMessage(content=\"Do not hallucinate or assume any information. For example, if annual expenses are not provided, ask the user for it. After getting annual expenses from the user, call the tool 'calculate_fire_number' to get the FIRE number for the user.\", additional_kwargs={}, response_metadata={}), HumanMessage(content='What is my FIRE number? 10K.', additional_kwargs={}, response_metadata={}), ToolMessage(content='250000.0', name='calculate_fire_number', tool_call_id='1010a84e-5550-4a80-89cd-ff36b15c60d7')]\n"
     ]
    }
   ],
   "source": [
    "name_tool_mapping = {\"calculate_fire_number\" : calculate_fire_number}\n",
    "\n",
    "for tool_call in response.tool_calls:\n",
    "    selected_tool = name_tool_mapping[tool_call[\"name\"]]\n",
    "    messages.append(selected_tool.invoke(tool_call))\n",
    "\n",
    "print(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "content='The FIRE number was calculated by multiplying your expected annual expenses of 10000 by 25. So, the calculation is:\\n\\n\\\\[ \\\\text{FIRE number} = 10000 \\\\times 25 = 250000 \\\\] \\n\\nThis means you would need savings of at least 250000 to achieve financial independence based on your spending habits.' additional_kwargs={} response_metadata={'model': 'qwen2.5:7b', 'created_at': '2025-02-16T08:51:08.279939Z', 'done': True, 'done_reason': 'stop', 'total_duration': 12892559292, 'load_duration': 28216459, 'prompt_eval_count': 258, 'prompt_eval_duration': 4697000000, 'eval_count': 87, 'eval_duration': 8158000000, 'message': Message(role='assistant', content='The FIRE number was calculated by multiplying your expected annual expenses of 10000 by 25. So, the calculation is:\\n\\n\\\\[ \\\\text{FIRE number} = 10000 \\\\times 25 = 250000 \\\\] \\n\\nThis means you would need savings of at least 250000 to achieve financial independence based on your spending habits.', images=None, tool_calls=None)} id='run-6d057c0b-31ce-4ae1-a12c-85fa112c8e03-0' usage_metadata={'input_tokens': 258, 'output_tokens': 87, 'total_tokens': 345}\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "\n",
    "system_template = \"The user had asked to calculate his FIRE number. The user had expected annual expenses of {annual_expenses}. We already ran a tool to calculate the FIRE number, the output of which was {fire_number}. Please share this response with the user, telling them how this number was calculated (by multiplying the annual expenses by 25). Do not assume any other information apart from what has been shared in this context. Do not format the FIRE number with any currency.\"\n",
    "\n",
    "prompt_template = ChatPromptTemplate.from_messages([(\"system\", system_template)])\n",
    "\n",
    "prompt = prompt_template.invoke({\"annual_expenses\": response.tool_calls[0]['args']['annual_expenses'], \"fire_number\": messages[-1].content})\n",
    "\n",
    "# print(prompt)\n",
    "\n",
    "final_response = model_with_tools.invoke(prompt)\n",
    "\n",
    "print(final_response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
